#!/bin/bash

# 运行测试的脚本

# 设置颜色
RED='\033[0;31m'
GREEN='\033[0;32m'
YELLOW='\033[0;33m'
BLUE='\033[0;34m'
NC='\033[0m' # 无颜色

# 检查是否存在虚拟环境
if [ ! -d ".venv" ]; then
    echo -e "${YELLOW}警告: 未找到虚拟环境，正在创建...${NC}"
    python3 -m venv .venv
    if [ $? -ne 0 ]; then
        echo -e "${RED}创建虚拟环境失败${NC}"
        exit 1
    fi
    echo -e "${GREEN}虚拟环境已创建${NC}"
fi

# 激活虚拟环境
echo -e "${BLUE}激活虚拟环境...${NC}"
echo "当前目录: $(pwd)"
echo "虚拟环境路径: $(pwd)/.venv"
if [ -f .venv/bin/activate ]; then
    source .venv/bin/activate
    if [ $? -ne 0 ]
    then
        echo -e "${RED}激活虚拟环境失败${NC}"
        exit 1
    fi
    echo -e "${GREEN}虚拟环境已激活${NC}"
    echo "激活后Python路径: $(which python3)"
    echo "激活后Python版本: $(python3 --version)"
    echo "激活后pip路径: $(which pip3)"
    echo "激活后pip版本: $(pip3 --version)"
else
    echo -e "${RED}错误: 虚拟环境激活脚本不存在${NC}"
    exit 1
fi

# 检查是否安装了Python
if ! command -v python3 &> /dev/null

then
    echo -e "${RED}错误: 未找到Python3，请先安装Python3${NC}"
    exit 1
fi

# 检查是否安装了pip
if ! command -v pip3 &> /dev/null
then
    echo -e "${RED}错误: 未找到pip3，请先安装pip3${NC}"
    exit 1
fi

# 验证langchain_community是否安装
echo -e "${BLUE}验证langchain_community是否安装...${NC}"
if python3 -c "import langchain_community" &> /dev/null
then
    echo -e "${GREEN}langchain_community已安装${NC}"
else
    echo -e "${RED}错误: langchain_community未安装${NC}"
    pip3 install langchain-community
    if [ $? -ne 0 ]
    then
        echo -e "${RED}安装langchain_community失败${NC}"
        exit 1
    fi
fi

# 检查sentence-transformers是否安装
echo -e "${BLUE}检查sentence-transformers是否安装...${NC}"
if python3 -c "import sentence_transformers" &> /dev/null
then
    echo -e "${GREEN}sentence-transformers已安装${NC}"
else
    echo -e "${YELLOW}警告: sentence-transformers未安装，尝试安装...${NC}"
    pip3 install sentence-transformers>=2.2.2
    if [ $? -ne 0 ]
    then
        echo -e "${RED}安装sentence-transformers失败${NC}"
        exit 1
    fi
fi

# 检查Milvus服务是否运行
echo -e "${BLUE}检查Milvus服务是否运行...${NC}"
# 尝试连接Milvus服务
if ! nc -z localhost 19530
 then
    echo -e "${YELLOW}警告: Milvus服务未运行，请先启动Milvus服务${NC}"
    echo -e "${BLUE}Milvus启动指南:${NC}"
    echo -e "1. 确保已安装Docker和Docker Compose"
    echo -e "2. 访问Milvus官方文档: https://milvus.io/docs/install_standalone-docker.md"
    echo -e "3. 按照文档说明启动Milvus服务"
    echo -e "${YELLOW}启动Milvus服务后，重新运行此脚本${NC}"
    exit 1
fi

echo -e "${GREEN}Milvus服务已运行${NC}"

# 检查嵌入模型和rerank模型路径
 echo -e "${BLUE}检查模型路径...${NC}"
 EMBEDDING_MODEL_PATH="/Users/apple/.cache/huggingface/hub/models--shibing624--text2vec-base-chinese/snapshots/183bb99aa7af74355fb58d16edf8c13ae7c5433e"
 RERANK_MODEL_PATH="/Users/apple/.cache/huggingface/hub/models--cross-encoder--mmarco-mMiniLMv2-L12-H384-v1/snapshots/1427fd652930e4ba29e8149678df786c240d8825"

 if [ -d "$EMBEDDING_MODEL_PATH" ]
 then
     echo -e "${GREEN}嵌入模型路径存在: $EMBEDDING_MODEL_PATH${NC}"
 else
     echo -e "${RED}错误: 嵌入模型路径不存在: $EMBEDDING_MODEL_PATH${NC}"
     echo -e "${BLUE}请确保已下载模型或更新test_embedding.py中的模型路径${NC}"
     exit 1
 fi

 if [ -d "$RERANK_MODEL_PATH" ]
 then
     echo -e "${GREEN}rerank模型路径存在: $RERANK_MODEL_PATH${NC}"
 else
     echo -e "${YELLOW}警告: rerank模型路径不存在: $RERANK_MODEL_PATH${NC}"
     echo -e "${BLUE}将使用不启用重排序的模式运行测试${NC}"
 fi


# 运行测试
echo -e "${BLUE}开始运行测试...${NC}"
python3 test_doc_embedding.py
if [ $? -ne 0 ]
then
    echo -e "${RED}测试运行失败${NC}"
    exit 1
else
    echo -e "${GREEN}测试运行成功${NC}"
    echo -e "${BLUE}测试结果已保存到test_results.json${NC}"
fi

echo ""
echo ""
echo ""
echo "--------------------------------------------------"
python3 test_image_embedding.py
if [ $? -ne 0 ]
then
    echo -e "${RED}测试运行失败${NC}"
    exit 1
else
    echo -e "${GREEN}测试运行成功${NC}"
    echo -e "${BLUE}测试结果已保存到test_results.json${NC}"
fi

exit 0